import os
import torch
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import gradio as gr
from model import FruitClassifier
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')  # 使用非交互式后端
from colorama import Fore, Style
import time

# 设置中文字体
try:
    # 尝试设置微软雅黑字体
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'Arial Unicode MS', 'sans-serif']
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
except:
    print("警告: 无法设置中文字体，图表中的中文可能无法正确显示")

# 图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

def load_model(model_path, num_classes=30, device=None):
    """加载训练好的模型"""
    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    print(f"{Fore.CYAN}使用设备: {device}{Style.RESET_ALL}")
    
    try:
        # 加载模型
        print(f"{Fore.YELLOW}正在加载模型: {model_path}{Style.RESET_ALL}")
        checkpoint = torch.load(model_path, map_location=device)
        
        model = FruitClassifier(num_classes=num_classes)
        model.load_state_dict(checkpoint['model_state_dict'])
        model = model.to(device)
        model.eval()
        
        print(f"{Fore.GREEN}模型加载成功!{Style.RESET_ALL}")
        return model, device
    except Exception as e:
        print(f"{Fore.RED}模型加载失败: {str(e)}{Style.RESET_ALL}")
        return None, device

def get_class_names(data_dir):
    """获取类别名称"""
    if not os.path.exists(data_dir):
        print(f"{Fore.RED}数据目录不存在: {data_dir}{Style.RESET_ALL}")
        # 返回一些默认类别
        return ["苹果", "香蕉", "橙子", "草莓", "西瓜"]
    
    # 获取子目录名称作为类别
    class_names = sorted([d for d in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, d))])
    
    if not class_names:
        print(f"{Fore.RED}在数据目录中未找到类别: {data_dir}{Style.RESET_ALL}")
        # 返回一些默认类别
        return ["苹果", "香蕉", "橙子", "草莓", "西瓜"]
    
    return class_names

def predict_image(image, model, device, class_names):
    """预测图像的类别"""
    if image is None:
        return None, None, None
    
    # 转换图像格式
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image.astype('uint8'))
    
    # 应用预处理
    img_tensor = transform(image).unsqueeze(0).to(device)
    
    # 预测
    with torch.no_grad():
        outputs = model(img_tensor)
        probs = F.softmax(outputs, dim=1)[0]
        
    # 获取前5个预测结果
    top5_probs, top5_indices = torch.topk(probs, 5)
    top5_probs = top5_probs.cpu().numpy()
    top5_indices = top5_indices.cpu().numpy()
    
    # 获取类别名称
    top5_labels = [class_names[idx] for idx in top5_indices]
    
    # 创建条形图
    fig, ax = plt.subplots(figsize=(10, 6))
    y_pos = np.arange(len(top5_labels))
    
    # 使用英文标签避免中文显示问题
    display_labels = [f"Class {i}" for i in range(len(top5_labels))]
    
    # 创建水平条形图 - 确保概率值在0-100%之间
    bars = ax.barh(y_pos, top5_probs * 100, color='skyblue')
    ax.set_yticks(y_pos)
    ax.set_yticklabels(display_labels)
    ax.invert_yaxis()  # 标签从上到下
    ax.set_xlabel('Confidence (%)')
    ax.set_title('Top 5 Predictions')
    
    # 在条形上添加数值标签 - 确保格式正确
    for i, bar in enumerate(bars):
        width = bar.get_width()
        ax.text(width + 1, bar.get_y() + bar.get_height()/2, f'{width:.1f}%',
                ha='left', va='center', fontsize=10)
    
    # 添加类别映射说明
    plt.figtext(0.5, 0.01, "Class mapping: " + ", ".join([f"{i}: {label}" for i, label in enumerate(top5_labels)]), 
                ha="center", fontsize=8, wrap=True)
    
    # 设置x轴范围，留出空间给标签
    ax.set_xlim(0, 110)
    
    plt.tight_layout()
    
    # 保存图像到临时文件
    temp_file = f"temp_plot_{time.time()}.png"
    plt.savefig(temp_file)
    plt.close(fig)
    
    # 返回预测结果 - 确保概率值格式正确
    prediction = {
        "类别": top5_labels[0],
        "置信度": f"{top5_probs[0] * 100:.2f}%"
    }
    
    # 创建HTML表格显示详细结果
    html_content = "<table style='width:100%; border-collapse: collapse;'>"
    html_content += "<tr><th style='text-align:left; padding:8px; border-bottom:1px solid #ddd;'>类别</th><th style='text-align:right; padding:8px; border-bottom:1px solid #ddd;'>置信度</th></tr>"
    
    for label, prob in zip(top5_labels, top5_probs):
        html_content += f"<tr><td style='text-align:left; padding:8px; border-bottom:1px solid #ddd;'>{label}</td>"
        html_content += f"<td style='text-align:right; padding:8px; border-bottom:1px solid #ddd;'>{prob * 100:.2f}%</td></tr>"
    
    html_content += "</table>"
    
    # 返回图像路径而不是图像对象
    return prediction, temp_file, html_content

def create_ui(model, device, class_names):
    """创建Gradio界面"""
    
    # 定义界面布局
    with gr.Blocks(title="水果分类系统", theme=gr.themes.Soft()) as demo:
        gr.Markdown(
            """
            # 水果分类系统
            
            上传一张水果图片，AI将识别它属于哪种水果类别。
            
            ## 使用说明
            1. 点击下方上传图片或拖放图片到上传区域
            2. 等待AI分析结果
            3. 查看预测结果和置信度
            """
        )
        
        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(type="pil", label="上传水果图片")
                with gr.Row():
                    submit_btn = gr.Button("开始识别", variant="primary")
                    clear_btn = gr.Button("清除", variant="secondary")
                
            with gr.Column(scale=1):
                with gr.Group():
                    output_label = gr.JSON(label="预测结果")
                    # 使用Image组件代替Plot组件
                    output_confidence = gr.Image(label="置信度分布", type="filepath")
                    # 使用HTML组件代替Label组件，以便自定义格式
                    output_table = gr.HTML(label="详细结果")
        
        # 设置事件
        submit_btn.click(
            fn=lambda img: predict_image(img, model, device, class_names),
            inputs=input_image,
            outputs=[output_label, output_confidence, output_table]
        )
        
        clear_btn.click(
            fn=lambda: (None, None, None),
            inputs=None,
            outputs=[input_image, output_confidence, output_table]
        )
        
        # 添加示例
        gr.Examples(
            examples=[
                os.path.join("examples", "example1.jpg") if os.path.exists(os.path.join("examples", "example1.jpg")) else None,
                os.path.join("examples", "example2.jpg") if os.path.exists(os.path.join("examples", "example2.jpg")) else None,
                os.path.join("examples", "example3.jpg") if os.path.exists(os.path.join("examples", "example3.jpg")) else None,
            ],
            inputs=input_image,
            outputs=[output_label, output_confidence, output_table],
            fn=lambda img: predict_image(img, model, device, class_names),
            cache_examples=True,
        )
    
    return demo

def main():
    # 模型路径
    model_path = "model/best_model.pth"
    
    # 数据目录（用于获取类别名称）
    data_dir = "plant_data/Train_Set_Folder"
    
    # 获取类别名称
    class_names = get_class_names(data_dir)
    
    # 加载模型
    model, device = load_model(model_path, num_classes=len(class_names))
    
    if model is None:
        print(f"{Fore.RED}无法启动Web界面，模型加载失败{Style.RESET_ALL}")
        return
    
    # 创建示例目录
    if not os.path.exists("examples"):
        os.makedirs("examples")
        print(f"{Fore.YELLOW}创建了examples目录，你可以在这里放置示例图片{Style.RESET_ALL}")
    
    # 创建并启动UI
    demo = create_ui(model, device, class_names)
    
    # 启动Gradio应用
    demo.launch(share=True, server_name="127.0.0.1")

if __name__ == "__main__":
    main()